/* Copyright (C) 2009, 2011 Christofer Jonsson
 *
 * This file is part of FXH.
 *
 * FXH is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * FXH is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
 * License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with FXH.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "fxh_socket.h"
#include "fxh_error.h"
#include "fxh_crypto.h"
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <netinet/in.h>
#include <netdb.h>
#include <arpa/inet.h>

/* Bitmap */
static uint32_t options = 0;

static fxh_status_t send_all(fxh_socket_t *socket,
                             const void *buffer,
                             size_t size);

static fxh_status_t recv_all(fxh_socket_t *socket,
                             void *buffer,
                             size_t size);

static fxh_status_t send_buffer(fxh_socket_t *socket,
                                const void *buffer,
                                size_t size,
                                size_t *sent);

static fxh_status_t recv_buffer(fxh_socket_t *socket,
                                void *buffer,
                                size_t size,
                                size_t *received);

static void stats_update_last_trans(fxh_socket_t *socket,
				    size_t bytes,
				    struct timeval *start,
				    struct timeval *end);
				    

void fxh_socket_enable(uint32_t options)
{
   options |= options;
}

void fxh_socket_disable(uint32_t options)
{
   options &= ~(options);
}

fxh_status_t fxh_socket_listen(fxh_socket_t *listener,
                               const char *port,
                               int backlog)
{
   struct addrinfo hints, *addrs, *addr;
   int result, yes = 1;

   memset(&hints, 0, sizeof(hints));
   hints.ai_family = AF_UNSPEC;
   hints.ai_socktype = SOCK_STREAM;
   hints.ai_flags = AI_PASSIVE;

   /* Get address info */
   if ((result = getaddrinfo(NULL, port, &hints, &addrs)) != 0)
   {
      FXH_ERROR("Unable to get address info: %s", gai_strerror(result));
      return FXH_FAILURE;
   }

   for(addr = addrs; addr != NULL; addr = addr->ai_next)
   {
      /* Create socket */
      listener->fd = socket(addr->ai_family,
                            addr->ai_socktype,
                            addr->ai_protocol);

      if (listener->fd == -1)
      {
         FXH_ERROR("Unable to create socket: %s", strerror(errno));
         continue;
      }

      /* Set socket configuration */
      result = setsockopt(listener->fd,
                          SOL_SOCKET,
                          SO_REUSEADDR,
                          &yes, (socklen_t)sizeof(int));

      if (result == -1)
      {
         (void)close(listener->fd);
         FXH_ERROR("Unable to set SO_REUSEADDR to true: %s", strerror(errno));
         continue;
      }

      /* Bind */
      result = bind(listener->fd, addr->ai_addr, addr->ai_addrlen);

      if (result == -1)
      {
         (void)close(listener->fd);
         FXH_ERROR("Unable to bind: %s", strerror(errno));
         continue;
      }

      break;
   }

   freeaddrinfo(addrs);

   if (addr == NULL)
   {
      FXH_ERROR("Failed to bind");
      return FXH_FAILURE;
   }

   /* Listen */
   if (listen(listener->fd, backlog) == -1)
   {
      FXH_ERROR("Failed to listen");
      return FXH_FAILURE;
   }

   return FXH_SUCCESS;
}

fxh_status_t fxh_socket_connect(fxh_socket_t *connector,
                                const char *address,
                                const char *port)
{
   struct addrinfo hints, *addrs, *addr;
   int result;

   memset(&hints, 0, sizeof(hints));
   hints.ai_family = AF_UNSPEC;
   hints.ai_socktype = SOCK_STREAM;

   /* Get address info */
   if ((result = getaddrinfo(address, port, &hints, &addrs)) != 0)
   {
      FXH_ERROR("Unable to get address info: %s", gai_strerror(result));
      return FXH_FAILURE;
   }

   for(addr = addrs; addr != NULL; addr = addrs->ai_next)
   {
      /* Create socket */
      connector->fd = socket(addr->ai_family,
                             addr->ai_socktype,
                             addr->ai_protocol);

      if (connector->fd == -1)
      {
         FXH_ERROR("Unable to create socket: %s", strerror(errno));
         continue;
      }

      /* Connect */
      result = connect(connector->fd, addr->ai_addr, addr->ai_addrlen);

      if (result == -1)
      {
         (void)close(connector->fd);
         FXH_ERROR("Unable to connect: %s", strerror(errno));
         continue;
      }

      break;
   }

   freeaddrinfo(addrs);

   if (addr == NULL)
   {
      FXH_ERROR("Failed to connect");
      return FXH_FAILURE;
   }

   return FXH_SUCCESS;
}

fxh_status_t fxh_socket_accept(fxh_socket_t *listener,
                               fxh_socket_t *connector)
{
   struct sockaddr_storage addr;
   socklen_t size = (socklen_t)sizeof(addr);

   /* Accept */
   connector->fd = accept(listener->fd,
                          (struct sockaddr *)&addr, &size);

   if (connector->fd == -1)
   {
      FXH_ERROR("Unable to accept connecton: %s", strerror(errno));
      return FXH_FAILURE;
   }

   return FXH_SUCCESS;
}

fxh_status_t fxh_socket_send_all(fxh_socket_t *socket,
                                 void *buffer,
                                 size_t size)
{
   fxh_status_t result = FXH_FAILURE;
   struct timeval start, end;

   if (options & FXH_SOCKET_ENCRYPTION)
      fxh_crypto_encrypt(buffer, size);

   if (0 != gettimeofday(&start, NULL))
   {
      FXH_ERROR("Call to gettimeofday failed: %s", strerror(errno));
      return FXH_FAILURE;
   }

   result = send_all(socket, buffer, size);

   if (0 != gettimeofday(&end, NULL))
   {
      FXH_ERROR("Call to gettimeofday failed: %s", strerror(errno));
      return FXH_FAILURE;
   }

   stats_update_last_trans(socket, size, &start, &end);
  
   if (options & FXH_SOCKET_ENCRYPTION)
      fxh_crypto_decrypt(buffer, size);

   return result;
}

fxh_status_t fxh_socket_recv_all(fxh_socket_t *socket,
                                 void *buffer,
                                 size_t size)
{
   struct timeval start, end;
   fxh_status_t result = FXH_FAILURE;

   if (0 != gettimeofday(&start, NULL))
   {
      FXH_ERROR("Call to gettimeofday failed: %s", strerror(errno));
      return FXH_FAILURE;
   }

   result = recv_all(socket, buffer, size);

   if (0 != gettimeofday(&end, NULL))
   {
      FXH_ERROR("Call to gettimeofday failed: %s", strerror(errno));
      return FXH_FAILURE;
   }

   stats_update_last_trans(socket, size, &start, &end);

   if (options & FXH_SOCKET_ENCRYPTION)
      fxh_crypto_decrypt(buffer, size);

   return result;
}

fxh_status_t fxh_socket_send(fxh_socket_t *socket,
                             void *buffer,
                             size_t size,
                             size_t *sent)
{
   struct timeval start, end;
   fxh_status_t result = FXH_FAILURE;

   if (options & FXH_SOCKET_ENCRYPTION)
      fxh_crypto_encrypt(buffer, size);

   if (0 != gettimeofday(&start, NULL))
   {
      FXH_ERROR("Call to gettimeofday failed: %s", strerror(errno));
      return FXH_FAILURE;
   }

   result = send_buffer(socket, buffer, size, sent);

   if (0 != gettimeofday(&end, NULL))
   {
      FXH_ERROR("Call to gettimeofday failed: %s", strerror(errno));
      return FXH_FAILURE;
   }

   stats_update_last_trans(socket, *sent, &start, &end);

   if (options & FXH_SOCKET_ENCRYPTION)
      fxh_crypto_decrypt(buffer, size);

   return result;
}

fxh_status_t fxh_socket_recv(fxh_socket_t *socket,
                             void *buffer,
                             size_t size,
                             size_t *received)
{
   struct timeval start, end;
   fxh_status_t result = FXH_FAILURE;

   if (0 != gettimeofday(&start, NULL))
   {
      FXH_ERROR("Call to gettimeofday failed: %s", strerror(errno));
      return FXH_FAILURE;
   }

   result = recv_buffer(socket, buffer, size, received);

   if (0 != gettimeofday(&end, NULL))
   {
      FXH_ERROR("Call to gettimeofday failed: %s", strerror(errno));
      return FXH_FAILURE;
   }

   stats_update_last_trans(socket, *received, &start, &end);

   if (options & FXH_SOCKET_ENCRYPTION)
      fxh_crypto_decrypt(buffer, *received);

   return result;
}

fxh_status_t fxh_socket_close(fxh_socket_t *socket)
{
   if (close(socket->fd) == -1)
   {
      FXH_ERROR("Failed to close socket: %s", strerror(errno));
      return FXH_FAILURE;
   }

   return FXH_SUCCESS;
}

static fxh_status_t send_all(fxh_socket_t *socket,
                             const void *buffer,
                             size_t size)
{
   ssize_t bytes = 0;
   size_t total = 0;
   size_t left = size;

   while (total != size)
   {
      bytes = send(socket->fd, (char*)buffer + total, left, 0);

      if (bytes == -1)
      {
         FXH_ERROR("Unable to send data: %s", strerror(errno));
         return FXH_FAILURE;
      }

      total += bytes;
      left -= bytes;
   }

   return FXH_SUCCESS;
}

static fxh_status_t recv_all(fxh_socket_t *socket,
                             void *buffer,
                             size_t size)
{
   ssize_t bytes = 0;
   size_t total = 0;
   size_t left = size;

   while (total != size)
   {
      bytes = recv(socket->fd, (char*)buffer + total, left, 0);

      if (bytes == -1)
      {
         FXH_ERROR("Unable to receive data: %s", strerror(errno));
         return FXH_FAILURE;
      }

      if (bytes == 0)
      {
         FXH_ERROR("Connection closed by peer");
         return FXH_FAILURE;
      }

      total += bytes;
      left -= bytes;
   }

   return FXH_SUCCESS;
}

static fxh_status_t send_buffer(fxh_socket_t *socket,
                                const void *buffer,
                                size_t size,
                                size_t *sent)
{
   ssize_t bytes = send(socket->fd, buffer, size, 0);

   if (bytes == -1)
   {
      FXH_ERROR("Unable to send data: %s", strerror(errno));
      return FXH_FAILURE;
   }

   *sent = (size_t)bytes;
   return FXH_SUCCESS;
}

static fxh_status_t recv_buffer(fxh_socket_t *socket,
                                void *buffer,
                                size_t size,
                                size_t *received)
{
   ssize_t bytes = recv(socket->fd, buffer, size, 0);

   if (bytes == -1)
   {
      FXH_ERROR("Unable to receive data: %s", strerror(errno));
      return FXH_FAILURE;
   }

   if (bytes == 0)
   {
      FXH_ERROR("Connection closed by peer");
      return FXH_FAILURE;
   }

   *received = (size_t)bytes;
   return FXH_SUCCESS;
}

static void stats_update_last_trans(fxh_socket_t *socket,
				    size_t bytes,
				    struct timeval *start,
				    struct timeval *end)
{
   double start_ms = (start->tv_sec * 1000.0) + (start->tv_usec / 1000.0);
   double end_ms = (end->tv_sec * 1000.0) + (end->tv_usec / 1000.0);
   socket->stats.last_trans.bytes = bytes;
   socket->stats.last_trans.time = end_ms - start_ms;
}
